# visual_verify_vkitti.py

import os
import os.path as osp
import sys
import glob
import logging
import argparse
import numpy as np
from tqdm import tqdm
import cv2
# poses correct
# --- Import from our visual utility file ---
from visual_util import launch_visualizer

# --- Import from the original utility file for data loading helpers ---
try:
    from verify_utils import setup_logging, depthmap_to_absolute_camera_coordinates, imread_cv2
except ImportError:
    print("FATAL ERROR: Could not import from 'verify_utils.py'.")
    print("Please ensure 'verify_utils.py' is in the same directory.")
    sys.exit(1)


def load_vkitti_view_data(rgb_path, logger):
    """
    Loads all necessary data for a single view from the preprocessed vKITTI dataset format.
    """
    try:
        # Construct paths based on the vKITTI naming convention
        base_name = osp.basename(rgb_path).replace("_rgb.jpg", "")
        depth_path = osp.join(osp.dirname(rgb_path), f"{base_name}_depth.png")
        cam_path = osp.join(osp.dirname(rgb_path), f"{base_name}_cam.npz")

        if not all(osp.exists(p) for p in [depth_path, cam_path]):
            logger.warning(f"Missing files for {base_name}. Skipping.")
            return None

        # Load image and depth data
        rgb_image = imread_cv2(rgb_path) # Assumes it returns RGB
        depthmap = imread_cv2(depth_path, -1) # Use -1 for cv2.IMREAD_UNCHANGED
        
        if rgb_image is None or depthmap is None:
            logger.warning(f"Failed to read image or depth for {rgb_path}. Skipping.")
            return None

        # --- Process data according to vKITTI format ---
        # 1. Normalize RGB image for the viewer
        rgb_image = rgb_image.astype(np.float32) / 255.0
        # 2. Scale depth from centimeters to meters
        depthmap = depthmap.astype(np.float32) / 100.0
        depthmap[~np.isfinite(depthmap)] = 0

        # Load camera data
        camera_data = np.load(cam_path)
        intrinsics = camera_data['camera_intrinsics'].astype(np.float32)
        camera_pose = camera_data['camera_pose'].astype(np.float32)

        return {
            "rgb": rgb_image, "depth": depthmap, "intrinsics": intrinsics,
            "pose": camera_pose, "path": rgb_path
        }
    except Exception as e:
        logger.error(f"Could not process view for {osp.basename(rgb_path)}. Details: {e}", exc_info=True)
        return None

# ======================================================================================
# --- Main Execution Block ---
# ======================================================================================

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Visually verify the consistency of a PREPROCESSED vKITTI dataset."
    )
    
    parser.add_argument(
        "--data_root", 
        type=str,
        default='/lc/data/3D/vkitti',
        help="Path to the root directory of the vKITTI dataset."
    )
    
    parser.add_argument(
        "--log_file", 
        type=str, 
        default="logs/visual_verification_vkitti.log",
        help="Path to the output log file."
    )

    parser.add_argument(
        "--scene", 
        type=str, 
        default='Scene18',
        help="Optional: Specify a single scene name (e.g., 'Scene01') to visualize."
    )

    parser.add_argument(
        "--seq", 
        type=str, 
        default='30-deg-left',
        help="Optional: Specify a single sequence (e.g., '0001') to visualize."
    )

    parser.add_argument(
        "--cam", 
        type=str, 
        default=None,
        choices=["Camera_0", "Camera_1"],
        help="Optional: Specify a single camera ('Camera_0' or 'Camera_1') to visualize."
    )

    parser.add_argument(
        "--port", 
        type=int, 
        default=8080,
        help="Port to run the Viser web server on."
    )

    parser.add_argument(
        "--split",
        type=str,
        default="all",
        choices=["train", "test", "all"],
        help="Which split of the data to use."
    )
    
    args = parser.parse_args()
    
    log_dir = osp.dirname(args.log_file)
    if log_dir:
        os.makedirs(log_dir, exist_ok=True)
    logger = setup_logging(args.log_file)

    if not osp.isdir(args.data_root):
        logger.error(f"Provided data root '{args.data_root}' does not exist.")
        sys.exit(1)
    
    # --- Scene Discovery ---
    if args.scene:
        top_level_scenes = [args.scene]
    else:
        top_level_scenes = sorted([d for d in os.listdir(args.data_root) if osp.isdir(osp.join(args.data_root, d))])

    if args.split == "train":
        top_level_scenes = top_level_scenes[:-1]
    elif args.split == "test":
        top_level_scenes = top_level_scenes[-1:]
    
    logger.info(f"Found {len(top_level_scenes)} top-level scenes for split '{args.split}'.")
    
    for scene in top_level_scenes:
        scene_path = osp.join(args.data_root, scene)
        
        if args.seq:
            seq_paths = [osp.join(scene_path, args.seq)]
            if not osp.isdir(seq_paths[0]):
                logger.error(f"Specified sequence directory '{seq_paths[0]}' does not exist. Skipping.")
                continue
        else:
            seq_paths = sorted([osp.join(scene_path, d) for d in os.listdir(scene_path) if osp.isdir(osp.join(scene_path, d))])
        
        for seq_path in seq_paths:
            if args.cam:
                cams_to_process = [args.cam]
            else:
                cams_to_process = ["Camera_0", "Camera_1"]

            for cam in cams_to_process:
                subscene_path = osp.join(seq_path, cam)
                if not osp.isdir(subscene_path):
                    logger.warning(f"Camera directory not found: {subscene_path}. Skipping.")
                    continue

                scene_key = f"{scene}/{osp.basename(seq_path)}/{cam}"
                logger.info(f"\n--- Processing subscene: {scene_key} ---")

                rgb_paths = sorted(glob.glob(osp.join(subscene_path, "*_rgb.jpg")))
                if not rgb_paths:
                    logger.warning(f"No RGB images found in {subscene_path}. Skipping.")
                    continue

                pc_list, color_list, conf_list = [], [], []
                focals, pps, Rs, ts = [], [], [], []
                
                logger.info("Loading views and generating point clouds...")
                for rgb_path in tqdm(rgb_paths, file=sys.stdout):
                    view_data = load_vkitti_view_data(rgb_path, logger)
                    if view_data:
                        world_points, valid_mask, _ = depthmap_to_absolute_camera_coordinates(
                            view_data["depth"], view_data["intrinsics"], view_data["pose"]
                        )
                        pc_list.append(world_points)
                        color_list.append(view_data['rgb'])
                        conf_list.append(valid_mask.astype(np.float32))

                        intrinsics = view_data['intrinsics']
                        pose = view_data['pose']
                        focals.append(intrinsics[0, 0])
                        pps.append((intrinsics[0, 2], intrinsics[1, 2]))
                        Rs.append(pose[:3, :3])
                        ts.append(pose[:3, 3])
            
                if not pc_list:
                    logger.warning(f"No valid data could be loaded for {scene_key}. Skipping.")
                    continue

                cam_dict = {'focal': focals, 'pp': pps, 'R': Rs, 't': ts}

                try:
                    logger.info(f"Launching interactive visualizer for {scene_key} at http://localhost:{args.port}")
                    logger.info("--> Close the browser tab or press Ctrl+C in the console to continue.")
                    launch_visualizer(
                        pc_list=pc_list, color_list=color_list, conf_list=conf_list,
                        cam_dict=cam_dict, port=args.port
                    )
                except KeyboardInterrupt:
                    logger.info(f"\nVisualizer for {scene_key} closed, proceeding...")
                except Exception as e:
                    logger.error(f"Failed to launch visualizer for {scene_key}: {e}", exc_info=True)
                break

    logger.info("\n--- Visual verification complete. ---")
